"""
主函数，Text2SQL 用于将用户需求 转换为 SQL查询

t2s = Text2SQL(...)
sql = t2s.generate(query=query)

"""

# 导入数据源
from nl2sql.datasource.base import BaseDatasource
from nl2sql.datasource.bird_datasource import SqliteDataSource

# 导入必要的组件，如模型、向量模型等
from nl2sql.model.llm import LLM
from nl2sql.model.embedding_model import EmbeddingModel

# 导入我们创建的辅助Agent
from nl2sql.engines.condition_parser import ConditionParser
from nl2sql.engines.pseudo_code import PseudoCodeGenerator
from nl2sql.engines.text2sql.text2sql import SQLGeneratorAgent


class Text2SQL:
    def __init__(self,
                 llm: LLM,
                 embedding_model: EmbeddingModel = None,
                 datasource: BaseDatasource = None
                 ):
        self.llm: LLM = llm
        self.datasource = datasource

        if self.llm is None:
            raise ValueError('llm is None')

        self.embedding_model: EmbeddingModel = embedding_model

        self.sql_agent = SQLGeneratorAgent(llm=self.llm)

    def add_datasource(self, datasource: BaseDatasource) -> None:
        """给Text2SQL实例添加数据源"""
        self.datasource = datasource
        return

    @staticmethod
    def extract_sql_code(text) -> str:
        """从字符串中提取第一个SQL代码块"""
        import re

        # 定义正则表达式模式，用于匹配 ```sql 和 ``` 之间的内容
        pattern = r'```sql\s*([\s\S]*?)\s*```'

        # 使用正则表达式查找第一个匹配项
        match = re.search(pattern, text)

        # 如果找到匹配项，返回SQL代码，否则返回空字符串
        return match.group(1).strip() if match else ''

    def generate_sql(self,
                     query: str,
                     evidence: str) -> str:
        """
        用于生成准确的SQL
        """
        # 获取到数据源的信息
        data_info: str = self.datasource.get_data_info()

        sql = self.sql_agent.text2sql(query=query,
                                      evidence=evidence,
                                      data_info=data_info)

        return self.extract_sql_code(sql)


if __name__ == '__main__':
    from nl2sql.model.llm import LLM

    api_key: str = 'sk-68ac5f5ccf3540ba834deeeaecb48987'
    base_url: str = "https://dashscope.aliyuncs.com/compatible-mode/v1"
    model_name: str = "deepseek-v3"

    llm = LLM(api_key=api_key,
              base_url=base_url,
              model_name=model_name,
              temperature=0)

    sqlite_ds = SqliteDataSource(
        db_path='/Users/tiantiantian/bird_benchmark_dataset/dev/dev_20240627/dev_databases/california_schools/california_schools.sqlite',
        name='california_schools')

    t2s = Text2SQL(llm=llm, datasource=sqlite_ds)
    sql = t2s.generate_sql(
        query='What is the highest eligible free rate for K-12 students in the schools in Alameda County?',
        evidence='Eligible free rate for K-12 = `Free Meal Count (K-12)` / `Enrollment (K-12)`')

    sqlite_ds.execute_sql(sql=sql)
